import json
import os
from typing import List, Dict, Any
from rich.console import Console
from rich.panel import Panel
from rich.progress import Progress, SpinnerColumn, TextColumn
from rich.table import Table
from rich.markdown import Markdown
from embedding_api import EmbeddingAPI
from faiss_db import FaissDB

console = Console()

class WikiSearch:
    def __init__(self, data_dir: str = "wiki_chunks"):
        self.api = EmbeddingAPI()
        self.db = FaissDB()
        self.data_dir = data_dir

    def _load_chunks(self) -> List[Dict[str, Any]]:
        """加载所有分片数据"""
        chunks = []
        files = [f for f in os.listdir(self.data_dir) if f.endswith(".json")]
        
        with Progress(
            SpinnerColumn(),
            TextColumn("[progress.description]{task.description}"),
            transient=True
        ) as progress:
            task = progress.add_task("[cyan]加载分片数据...", total=len(files))
            for filename in files:
                try:
                    with open(os.path.join(self.data_dir, filename), 'r', encoding='utf-8') as f:
                        chunk = json.load(f)
                        if all(k in chunk for k in ["text", "context", "source_file", "char_length"]):
                            chunks.append(chunk)
                except Exception as e:
                    console.print(f"[red]Error loading {filename}: {str(e)}[/]")
                progress.update(task, advance=1)
        
        return chunks

    def build_index(self, test_mode: bool = False):
        """构建索引"""
        chunks = self._load_chunks()
        if not chunks:
            console.print("[red]未找到有效分片数据[/]")
            return

        if test_mode:
            chunks = chunks[:10]
            console.print(Panel.fit(f"[yellow]测试模式 - 仅处理前10个分片[/]"))

        # 分批处理
        batch_size = 16
        total = len(chunks)
        
        with Progress(
            SpinnerColumn(),
            TextColumn("[progress.description]{task.description}"),
            transient=True
        ) as progress:
            task = progress.add_task("[green]生成嵌入向量...", total=total)
            
            for i in range(0, total, batch_size):
                batch = chunks[i:i + batch_size]
                texts = [chunk["text"] for chunk in batch]
                
                # 使用独立的进度条显示当前批次状态
                with console.status(f"处理批次 {i//batch_size + 1}/{total//batch_size + 1}"):
                    embeddings = self.api.batch_embed(texts)
                    valid_data = [
                        (chunk, emb) 
                        for chunk, emb in zip(batch, embeddings) 
                        if emb is not None
                    ]
                    
                    if valid_data:
                        valid_chunks, valid_embs = zip(*valid_data)
                        self.db.add_chunks(valid_chunks, valid_embs)
                
                progress.update(task, advance=len(batch))

        self.db.save()
        console.print(Panel.fit("[green]✓ 索引构建完成[/]", border_style="green"))

    def interactive_search(self):
        """交互式搜索"""
        console.print(Panel.fit("""
        [bold blue]星穹铁道Wiki搜索系统[/]
        输入查询内容或输入 [yellow]/exit[/] 退出
        """, border_style="blue"))

        while True:
            query = console.input("\n[bold]🔍 搜索: [/]")
            if query.lower() in ('/exit', 'exit'):
                break

            with console.status("[cyan]搜索中...", spinner="dots"):
                embedding = self.api.get_embedding(query)
                if not embedding:
                    console.print("[red]无法生成查询向量[/]")
                    continue
                
                results = self.db.search(embedding)

            if not results:
                console.print("[yellow]未找到相关结果[/]")
                continue

            table = Table(title=f"搜索结果 (共{len(results)}条)", show_lines=True)
            table.add_column("相关度", justify="right")
            table.add_column("内容预览")
            table.add_column("来源文件")
            table.add_column("字符数")

            for result in results:
                table.add_row(
                    f"{result['score']:.3f}",
                    Markdown(result['text'][:500] + "..."),
                    os.path.basename(result['source_file']),
                    str(result['char_length'])
                )

            console.print(table)

if __name__ == "__main__":
    searcher = WikiSearch()
    
    # 构建索引（测试模式）
    # searcher.build_index(test_mode=False)
    
    # 交互式搜索
    searcher.interactive_search()
